load ./libgeoconstraint.so.0.0.0
package require GeoConstraint

# build GroupList 
proc buildGroupList { groupname targetAtomPdb selections args} {
    upvar $groupname groups
    global coords
    set targets {}
    set inStream [open $targetAtomPdb r]
    set targetPdbTop(segnames) {}
    set targetPdbTop(residues) {}
    set idx 0
    set groupby residue
    if {[set idx [lsearch $args by ]] >= 0} {
        set groupby [lindex $args [expr $idx +1]]
    }
    switch $groupby {
        segment {
            set group "segment=\$segname"
        }
        residue {
            set group "segment=\$segname,resid=\$resid"
        }
    }
    foreach line [split [read $inStream] \n] {
        set type [string trim [string range $line 0 5]]
        set name [string trim [string range $line 12 15]]
        set resid [string trim [string range $line 22 25]]
        set x [string trim [string range $line 27 37]]
        set y [string trim [string range $line 38 45]]
        set z [string trim [string range $line 46 53]]
        set beta [string trim [string range $line 60 65]]
        set occupancy [string trim [string range $line 54 59]]
        set segname [string trim [string range $line 72 75]]
        if { [string compare $type ATOM ] != 0}  { continue}
        set atomidx [atomid $segname $resid $name]
        set coords($atomidx) [list $x $y $z]
        set atomidx [atomid $segname $resid $name]
        set mygroup [subst $group]
        if ![info exists MyGroups($mygroup)] {
            set MyGroups($mygroup) {}
        }
        foreach selection $selections {
            if [string match $selection $name ] {
                lappend MyGroups($mygroup) $atomidx
            }
        }
    }
    foreach name [lsort [array names MyGroups]] {
        if {[llength $MyGroups($name)] > 0 } {
            lappend groups $MyGroups($name)
        }
    }
    close $inStream
    return
}
#load mass for each atom into global array mymasses. 
proc loadMasses { psf} {
  global mymasses
  set instream [open $psf r]
  set start 0
  foreach line [split [read $instream] \n] {
    if {$start> 0} {
        set segname [string trim [string range $line 8 13]]
        set resid [string trim [string range $line 13 18]]
        set name  [string trim [string range $line 23 28]]
        set atomidx [atomid $segname $resid $name]
        set mass [string trim [string range $line 48 60]]
        set charge [string trim [string range $line 33 47]]
        set mymasses($start) $mass
        incr start
        # until reached atom number 
        if {$start > $natom} { break }
          continue
    }
    #start after found NATOM 
     if {[regexp -- {(\d+)\s\!NATOM} $line matched natom] >0 } { set start 1;}
  }
}

# Setup Pore with 
proc setupPore { headgroup forceconstant r0 k0 {as Harmonic} args } {
   global headgroups waters tails mymasses coords geo fftx ffty fftz waterBoundary geo_in geo_out
   
    # head group COM &std 
    set zave 0.0
    set zstd 0.0
    foreach group $headgroups {
        set comsum {0.0 0.0 0.0 }
        set totalmass 0.0
        foreach atom $group {
           set tmp [vecscale $mymasses($atom) $coords($atom)]
           set comsum [vecadd $comsum $tmp]
           set totalmass [expr $totalmass +$mymasses($atom)]
           addatom $atom
        }
        set com [vecscale [expr 1.0/$totalmass] $comsum]
        set z [lindex $com 2]
        set zave [expr $zave+ abs($z)]
        set zstd [expr $zstd+ $z*$z]
    }
    set zave [expr $zave/[llength $headgroups]]
    set zstd [expr $zstd/[llength $headgroups]]
    set zstd [expr sqrt($zstd-$zave*$zave)]
    print "Head Group COM: $zave +/- $zstd"
    #get waters boundary
    set waterBoundary 999
    if ![info exists waters] {set waters {}}
    set idx [lsearch $args --constraint-water]
    if { $idx >= 0 && $idx < [expr [llength $args]-1] } {set forceconstant1 [lindex $args [expr $idx +1] ]} else {
     set waters {}
    }
    foreach group $waters {
        set comsum {0.0 0.0 0.0 }
        set totalmass 0.0
        foreach atom $group {
           set tmp [vecscale $mymasses($atom) $coords($atom)]
           set comsum [vecadd $comsum $tmp]
           set totalmass [expr $totalmass +$mymasses($atom)]
           addatom $atom
        }
        set com [vecscale [expr 1.0/$totalmass] $comsum]
        set z [expr abs([lindex $com 2])]
        if {$z < $waterBoundary } {
           set waterBoundary $z
        }
    }
    
    # add tails 
    if ![info exists waters] {set tails {}}
    set idx [lsearch $args --constraint-tails]
    if { $idx >= 0 && $idx < [expr [llength $args]-1] } {set forceconstant2 [lindex $args [expr $idx +1] ]} else {
      set tails {}
    }
    foreach group $tails {
        foreach atom $group {
            addatom $atom 
        }
    }
    if {[llength $waters] > 0 } {print "Water lower boundary: $waterBoundary"}
    # create geo constraint using above information.
    geoCreate geo ToroidalPore $r0 $k0 [expr $zave-$zstd]
    setGeoRange geo [expr -$fftx/2] [expr $fftx/2] [expr -$ffty/2] [expr $ffty/2] [expr -$fftz/2] [expr $fftz/2]
    print "Box size: x: [expr -$fftx/2] to [expr $fftx/2] y: [expr -$ffty/2] to [expr $ffty/2] z: [expr -$fftz/2] to [expr $fftz/2]"
    genGeoConstraint geo $as $forceconstant $zstd $zstd both 200
    # create inside constraint
    if {[llength $waters] > 0 } {
    geoCreate geo_in ToroidalPore $r0 $k0 [expr $zave-$zstd]
    setGeoRange geo_in [expr -$fftx/2] [expr $fftx/2] [expr -$ffty/2] [expr $ffty/2] [expr -$fftz/2] [expr $fftz/2]
    genGeoConstraint geo_in $as $forceconstant1 $zstd $zstd inside 200
    }
    # create outside constraint
    if {[llength $tails] > 0 } {
    geoCreate geo_out ToroidalPore $r0 $k0 [expr $zave-$zstd]
    setGeoRange geo_out [expr -$fftx/2] [expr $fftx/2] [expr -$ffty/2] [expr $ffty/2] [expr -$fftz/2] [expr $fftz/2]
    genGeoConstraint geo_out $as $forceconstant1 $zstd $zstd outside 200
    }
}

set waters {}
proc calcforces { } {
    global atoms numatoms mymasses namd COM forceconstant groups time TclForceUpdateFreq geo  
    global headgroups  myforce
    global waters tails geo_in geo_out
    #load coorinates
    loadcoords coords
    set time [getstep ]
    if { [expr $time % $TclForceUpdateFreq ] != 0 && [array exist myforce]} {
        foreach atom [array names myforce] {
           addforce $atom $myforce($atom)
        }
        return 
    }
    # update force every TclForceUpdateFreq steps.
    #groups of atom index orginized like that { { atom1 atom2 atom3} group2 } 
    set totalener 0.0
    #for head group apply Geo constraint 
    foreach group $headgroups { 
        set comsum {0.0 0.0 0.0}
        set totalmass 0.0
        foreach atom $group {
          set tmp [vecscale $mymasses($atom) $coords($atom)]
          set comsum [vecadd $comsum $tmp]
          set totalmass [expr $totalmass +$mymasses($atom)]
        }
        set com [vecscale [expr 1.0/$totalmass] $comsum]
        lassign $com x0 y0 z0
        # load Constraint value pre-calulated grids.
        set tmp [getGeoConstraint geo $x0 $y0 $z0 v]
        set ener [lindex $tmp 0] 
        set force [lrange $tmp 1 end]
        set totalener [expr $totalener + $ener]
        set force [vecscale -1.0 $force ]
        foreach atom $group {
            #separate the force to individual atoms
            #save force into array for later calculation.
            set myforce($atom) [vecscale [expr $mymasses($atom)/$totalmass] $force]
            addforce $atom $myforce($atom)
        }
    } 
  print "TS: $time Geo Constraint Energy: $totalener kcal/mol."
  set violatedwater {} 
  foreach group $waters {
        set comsum {0.0 0.0 0.0}
        set totalmass 0.0
        foreach atom $group {
             set tmp [vecscale $mymasses($atom) $coords($atom)]
             set comsum [vecadd $comsum $tmp]
             set totalmass [expr $totalmass +$mymasses($atom)]
        }
        set com [vecscale [expr 1.0/$totalmass] $comsum]
        lassign $com x0 y0 z0
        set tmp [getGeoConstraint geo_out $x0 $y0 $z0 v]
        set ener [lindex $tmp 0] 
        set force [lrange $tmp 1 end]
        set totalener [expr $totalener + $ener]
        set force [vecscale -1.0 $force ]
        #print "$group $x0 $y0 $z0"
        if { $ener >= 0.0 } {
           lappend violatedwater $group
           foreach atom $group {
            #separate the force to individual atoms
            #do not add the force for following timesteps
              set iforce [vecscale [expr $mymasses($atom)/$totalmass] $force]
              addforce $atom $iforce
           }
        } 
  }
  if {[llength $violatedwater] > 0 } { print "Push water $violatedwater back." }
 set violatedtails {} 
  foreach group $tails {
        set comsum {0.0 0.0 0.0}
        set totalmass 0.0
        foreach atom $group {
             set tmp [vecscale $mymasses($atom) $coords($atom)]
             set comsum [vecadd $comsum $tmp]
             set totalmass [expr $totalmass +$mymasses($atom)]
        }
        set com [vecscale [expr 1.0/$totalmass] $comsum]
        lassign $com x0 y0 z0
        set tmp [getGeoConstraint geo_in $x0 $y0 $z0 v]
        set ener [lindex $tmp 0] 
        set force [lrange $tmp 1 end]
        set totalener [expr $totalener + $ener]
        set force [vecscale -1.0 $force ]
        #print "$group $x0 $y0 $z0"
        if { $ener >= 0.0 } {
           lappend violatedtails $group
           foreach atom $group {
            #separate the force to individual atoms
            #do not add the force for following timesteps
              set iforce [vecscale [expr $mymasses($atom)/$totalmass] $force]
              addforce $atom $iforce
           }
        } 
  }
  if {[llength $violatedwater] > 0 } { print "Push tails $violatedtails back." }
}




